import sys
from concurrent.futures import ThreadPoolExecutor
import importlib
import requests
import glob
import os

response = requests.get('http://127.0.0.1:1025/v1/models')
data = response.json()
model_name = data['data'][0]['id'] if data['data'] else None
print(f"检测到mindie server中启动了{model_name}模型。")
try:
    model_module_path = f"models.{model_name}"
    model_module = importlib.import_module(model_module_path)
    model_infer = model_module.model_infer
except ModuleNotFoundError:
    print(f"[ERROR]： 暂不支持{model_name},请确认案例代码下存在{model_name}.py\r")
    current_dir = os.path.dirname(os.path.abspath(__file__))
    models_dir = os.path.join(current_dir, 'models')
    model_files = glob.glob(os.path.join(models_dir, '*'))
    model_names = [os.path.splitext(os.path.basename(file))[0] for file in model_files]
    print("当前支持的模型有：")
    for i, model_name in enumerate(model_names, start=1):
        print(f"{i}. {model_name}")

infer_engine = model_infer()

def inference_cli():
    while True:
        line = input()
        print(line)
        s=[i[0] for i in infer_engine.predict(line)]        # 非流式
        # s=[i[0] for i in infer_engine.predict_stream(line)] # 流式
        print(s)

        
def main():
    if len(sys.argv) > 1:
        inference_cli()
        return
    from flask import Flask, request, jsonify
    from flask import render_template  
    from flask_cors import CORS
    pool = ThreadPoolExecutor(max_workers=2)        
    app = Flask(
        __name__,
        static_folder='./dist',  
        template_folder="./dist",
        static_url_path=""
    )

    CORS(app, resources=r'/*')
    
    @app.route('/')
    def index():
        return render_template('index.html', name='index')

    @app.route("/api/chat", methods=["POST"])
    def getChat():
        msg = request.get_json(force=True)['message']
        if len(msg) == 0:
            return jsonify({"code": 404})
        print(msg)
        pool.submit(infer_engine.predict_stream,msg)
        return jsonify({"code": 200})

    @app.route("/api/getMsg", methods=["GET"])
    def getMsg():
        return jsonify(infer_engine.getState())
    
    @app.route("/api/reset", methods=["GET"])
    def reset():
        return jsonify({"code": 200})

    app.run(
        use_reloader=False,
        host="0.0.0.0",
        port=5000
    )

if __name__ == '__main__':
    main()

